import numpy as np
from math import sin, cos, radians, sqrt, degrees
import matplotlib.pyplot as plt

# --- I. 物理常数和弹丸参数 (基于提供文档) ---
# 国际标准大气模型参数
RHO0 = 1.225         # 海平面空气密度 (kg/m^3)
T0 = 288.15          # 海平面温度 (K)
L = 0.0065           # 温度递减率 (K/m)
G = 9.80665          # 重力加速度 (m/s^2)
R = 287.05           # 气体常数 J/(kg*K) 
P_TERM = (G / (R * L)) - 1 # 密度公式指数项

# 弹丸参数 (假设值)
C_SOUND = 338.0      # 海平面声速 (m/s)
M = 5.0              # 弹丸质量 (kg)
A = 0.0095           # 弹丸参考面积 (m^2)

# Cd-Ma 关系数据 (来自提供的曲线图)
MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5]) 
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40]) 

# --- II. 动态物理模型函数 ---

def get_drag_coefficient(v, c_sound):
    """根据马赫数线性插值计算 Cd。"""
    mach_number = v / c_sound
    Cd = np.interp(mach_number, MA_POINTS, CD_POINTS)
    return Cd

def get_air_density(z):
    """根据国际标准大气模型计算高度 z 处的空气密度 rho(z)。"""
    h = max(0.0, z) # 高度 h 即为 z
    temp_ratio = (T0 - L * h) / T0
    if temp_ratio <= 0:
        return 0.0 
    rho = RHO0 * (temp_ratio) ** P_TERM # 密度公式
    return rho

def derivatives(S, params):
    """
    计算状态向量 S = [x, z, vx, vz] 对时间 t 的导数 dS/dt。
    """
    x, z, vx, vz = S
    m, A, g, c_sound = params

    v = sqrt(vx**2 + vz**2) # 速度大小
    rho = get_air_density(z) # 动态 ρ(z)
    Cd = get_drag_coefficient(v, c_sound) # 动态 Cd(Ma) 
    
    drag_coeff = -(Cd * rho * A) / (2.0 * m) 
    
    # 加速度项 (dV/dt)
    dvx_dt = drag_coeff * v * vx 
    dvz_dt = drag_coeff * v * vz - g 
    
    # 导数向量 dS/dt = [dx/dt, dz/dt, dvx/dt, dvz/dt]
    return np.array([vx, vz, dvx_dt, dvz_dt])

def rk4_step(S, dt, params):
    """
    执行一步 RK4 积分。
    """
    K1 = dt * derivatives(S, params)
    K2 = dt * derivatives(S + K1 / 2.0, params)
    K3 = dt * derivatives(S + K2 / 2.0, params)
    K4 = dt * derivatives(S + K3, params)
    
    S_new = S + (K1 + 2.0 * K2 + 2.0 * K3 + K4) / 6.0
    return S_new

# --- III. 弹道模拟核心函数 (封装) ---

def calculate_trajectory(v0, theta_deg, dt=0.01):
    """
    使用 RK4 算法模拟弹道轨迹，并返回结果。
    """
    params = (M, A, G, C_SOUND)
    theta_rad = radians(theta_deg)
    
    # 初始状态 S = [x, z, vx, vz]
    S = np.array([
        0.0, 0.0, 
        v0 * cos(theta_rad), 
        v0 * sin(theta_rad)
    ])
    
    trajectory = []
    
    # 积分循环
    while S[1] >= -0.01: # 允许略微负值，确保接近地面
        # 记录当前状态
        trajectory.append(S.copy())
        
        S_new = rk4_step(S, dt, params)
        
        # 检查是否跨越地面
        if S_new[1] < 0 and S[1] >= 0:
            # 找到落地点
            # 这里可以进行更精细的插值，但为简化 RK4 演示，我们接受最后一个点作为落地点
            trajectory.append(S_new.copy())
            break
            
        S = S_new
    
    if not trajectory:
        return 0.0, 0.0, None

    traj_array = np.array(trajectory)
    R_final = traj_array[-1, 0] 
    Z_max = np.max(traj_array[:, 1])

    return R_final, Z_max, traj_array

# --- IV. 交互式运行和绘图 ---

def interactive_run():
    
    print("\n--- 🎯 RK4 弹道计算器 ---")
    print("请手动输入射角和初速度，程序将计算射程和最大射高。")
    
    results = []
    num_runs = 0
    
    # 循环接受用户输入
    while True:
        try:
            # 获取用户输入
            v0 = float(input("\n请输入初速度 v₀ (m/s)，输入 '0' 结束："))
            if v0 <= 0:
                break

            theta = float(input("请输入入射角 θ (度, 0-90)："))
            if not (0 <= theta <= 90):
                print("无效的射角。请重新输入。")
                continue
            
            # 计算轨迹
            R_final, Z_max, traj_data = calculate_trajectory(v0, theta)

            if R_final > 0:
                print("\n--- 结果 ---")
                print(f"✅ 初速度 v₀ = {v0:.3f} m/s, 射角 θ = {theta:.1f}°")
                print(f"   最终射程 R: {R_final:.3f} m")
                print(f"   最大射高 Z_max: {Z_max:.3f} m")
                
                # 存储结果用于绘图
                results.append({
                    'v0': v0, 
                    'theta': theta, 
                    'R': R_final, 
                    'data': traj_data
                })
                num_runs += 1
            else:
                print("❌ 积分失败或弹丸立即落地。请检查参数。")

        except ValueError:
            print("输入错误。请确保输入的是数字。")
        except Exception as e:
            print(f"发生错误: {e}")
    
    # 绘制所有轨迹
    if results:
        plot_trajectories(results)

def plot_trajectories(results):
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Heiti TC'] 
    plt.rcParams['axes.unicode_minus'] = False 
    plt.figure(figsize=(12, 7))
    
    max_range = 0
    max_height = 0
    
    # 绘制所有轨迹
    for i, res in enumerate(results):
        x_traj = res['data'][:, 0]
        z_traj = res['data'][:, 1]
        
        label = f'v₀={res["v0"]:.1f} m/s, θ={res["theta"]:.1f}°'
        
        plt.plot(x_traj, z_traj, label=label, linewidth=2, linestyle='-')
        
        # 更新图表范围
        max_range = max(max_range, res['R'])
        max_height = max(max_height, res['data'][:, 1].max())

    # 绘制地面
    plt.hlines(0, 0, max_range * 1.1, color='gray', linestyle='--', linewidth=0.8, label='地面')
    
    plt.title(f'RK4 弹道轨迹对比 (动态 $C_d(Ma), \\rho(z)$)', fontsize=16)
    plt.xlabel('水平距离 x (m)', fontsize=14)
    plt.ylabel('高度 z (m)', fontsize=14)
    plt.grid(True, linestyle=':', alpha=0.6)
    plt.legend(title='轨迹参数')
    plt.ylim(bottom=0)
    plt.xlim(right=max_range * 1.1)
    
    plt.show()

# 启动交互式运行
if __name__ == "__main__":
    interactive_run()